import pandas as pd
import numpy as np
import joblib
from sklearn.metrics import confusion_matrix, roc_auc_score, classification_report, roc_curve, accuracy_score
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from src.utils.common import data_processing
from sklearn.preprocessing import StandardScaler
from src.features.FeatureEngineering import feature_extra
from src.features.FeatureEngineering import feature_extra,EmployeeAttritionFeatureEngineer
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15


def model_predict():
    """
    1、模型预测
    2、模型评估
    :return: 返回真实值、阈值、预测值
    """
    x, y = data_processing(0)
    engineer = EmployeeAttritionFeatureEngineer()
    x = engineer.transform(x)
    # x, y = feature_extra(0)
    transfer = StandardScaler()
    x = transfer.fit_transform(x)

    es = joblib.load('../../model/logistic_model.pkl')
    y_predict = es.predict(x)
    y_scores = es.predict_proba(x)
    # print(y_scores)

    df_label = ['流失(正例)','留存(反例)']
    cm = confusion_matrix(1 - y, 1 - y_predict)
    df_cm = pd.DataFrame(cm, index=df_label, columns=df_label)
    print(f'预测模型对应的df结果：\n{df_cm}')

    print(f'流失AUC值为：{roc_auc_score(y, y_scores[:, 0])}')
    print(f'留存AUC值为：{roc_auc_score(y, y_scores[:, 1])}')
    print(f'roc为：\n{classification_report(y, y_predict, labels=[0, 1])}')
    print(f'准确率：{accuracy_score(y, y_predict)}')

    return y, y_scores, y_predict


def prediction_plot_ROC(true, scores):
    """
    绘制流失和留存ROC曲线
    :param true: 真实值
    :param scores: 阈值
    :return:
    """
    # 计算假正率、真正率、阈值
    fpr0, tpr0, thresholds0 = roc_curve(true, scores[:, 0])
    fpr1, tpr1, thresholds1 = roc_curve(true, scores[:, 1])
    # 计算AUC值
    roc_auc0 = roc_auc_score(true, scores[:, 0])
    roc_auc1 = roc_auc_score(true, scores[:, 1])

    # 绘制ROC曲线
    fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(20, 8), dpi=80)

    # 流失用户
    axes[0].plot(fpr0, tpr0, color='darkorange', lw=2, label=f'ROC曲线 (AUC = {roc_auc0})')
    # 对角线（随机猜测）
    axes[0].plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Random guess')
    # 刻度
    axes[0].set_xlim([0.0, 1.0])
    axes[0].set_ylim([0.0, 1.0])
    # 标题
    axes[0].set_xlabel('假正率 (FPR)')
    axes[0].set_ylabel('真正率 (TPR)')
    axes[0].set_title('流失用户特征曲线 (ROC)')
    # 图例
    axes[0].legend(loc="upper left")
    # 背景网格
    axes[0].grid(True, alpha=0.3)

    # 留存用户
    axes[1].plot(fpr1, tpr1, color='darkorange', lw=2, label=f'ROC曲线 (AUC = {roc_auc1})')
    axes[1].plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Random guess')
    axes[1].set_xlim([0.0, 1.0])
    axes[1].set_ylim([0.0, 1.0])
    axes[1].set_xlabel('假正率 (FPR)')
    axes[1].set_ylabel('真正率 (TPR)')
    axes[1].set_title('留存用户特征曲线 (ROC)')
    axes[1].legend(loc="lower right")
    axes[1].grid(True, alpha=0.3)

    plt.savefig('../../data/fig/流失和留存ROC.png')

    plt.show()


def predict_plot_confusion_matrix(true, predict):
    """
    混淆矩阵图示化
    :param true: 真实值
    :param predict: 预测值
    :return:
    """
    # 封装混淆矩阵为DataFrame
    df_label = ['流失(正例)', '留存(反例)']
    cm = confusion_matrix(1 - true, 1 - predict)
    df_cm = pd.DataFrame(cm, index=df_label, columns=df_label)

    # 创建画布
    fig, ax1 = plt.subplots(1, 1, figsize=(10, 10))

    # 颜色轴映射图形色块
    im = ax1.imshow([[0, 120], [150, 300]], cmap='viridis', vmin=0, vmax=300)

    # 刻度和标签
    ax1.set_xticks([0, 1], ['0/流失(正例)', '1/留存(反例)'])
    ax1.set_yticks([0, 1], ['0/流失(正例)', '1/留存(反例)'])
    ax1.set_yticklabels(['0/流失(正例)', '1/留存(反例)'], rotation=90, va='center')
    ax1.set_xlabel('预测标签')
    ax1.set_ylabel('真实标签')

    # 颜色轴基本设置
    axins = inset_axes(
        ax1,
        width="3%",  # width: 5% of parent_bbox width
        height="50%",  # height: 50%
        loc="lower right", # 颜色轴位置位与主图
        bbox_to_anchor=(1.05, 0., 1, 1), # 颜色轴锚点位置为（1.05， 0.），1.05：表示超出主图范围，在主图之外显示颜色轴
        bbox_transform=ax1.transAxes, # 表示颜色轴的比例是基于主图的比例进行生成的
        borderpad=0, # 颜色坐标轴与画布边界没有额外的填充空间
    )
    # 颜色轴刻度
    cbar = fig.colorbar(im, cax=axins)
    cbar.locator = plt.MaxNLocator(nbins=5)  # 自动选择最多 5 个刻度
    cbar.update_ticks()

    # 图形内各色块文本
    # j, i而不是i, j的原因是因为在DataFrame中先行索引（y）后列索引（x），但是画图时是先x后y
    # ha和va是水平和垂直居中
    # fontweight为加粗显示
    # np.indices(arr.shape)生成坐标网格，表示每个元素的位置；
    # np.vectorize(add_text)将add_text向量化，使其可对每个坐标位置调用；
    # *拆包坐标数组，作为add_text(i, j)的参数传入。
    def add_text(i, j):
        ax1.text(j, i, f'{df_cm.iloc[i, j]}', ha='center', va='center', color='red', fontsize=30, fontweight='bold')
    arr = np.array([[0, 120], [150, 300]])
    np.vectorize(add_text)(*np.indices(arr.shape))

    plt.savefig('../../data/fig/混淆矩阵.png')

    plt.show()


if __name__ == '__main__':
    true, scores, predict = model_predict()
    # prediction_plot_ROC(true, scores)
    # predict_plot_confusion_matrix(true, predict)
